#!/usr/bin/env python3
"""
捕获完整处理流程的HTTP响应
等待更长时间以确保获取到HTML生成结果
"""

import httpx
import asyncio
import sys
import subprocess
import time
import signal
import os

def cleanup_processes():
    """清理所有相关的进程"""
    print("[INFO] 正在清理相关进程...", file=sys.stderr)
    
    try:
        subprocess.run(["pkill", "-f", "uvicorn"], check=False, capture_output=True)
        subprocess.run(["pkill", "-f", "conversation_service"], check=False, capture_output=True)
        print("[INFO] 已清理相关进程", file=sys.stderr)
    except Exception as e:
        print(f"[WARN] 清理进程时出错: {e}", file=sys.stderr)
    
    # 检查端口8001是否被占用
    try:
        result = subprocess.run(["lsof", "-i", ":8001"], capture_output=True, text=True)
        if result.stdout:
            print("[INFO] 端口8001仍被占用，尝试强制清理...", file=sys.stderr)
            lines = result.stdout.strip().split('\n')[1:]
            for line in lines:
                if line:
                    parts = line.split()
                    if len(parts) > 1:
                        pid = parts[1]
                        try:
                            os.kill(int(pid), signal.SIGKILL)
                            print(f"[INFO] 已强制杀掉进程 {pid}", file=sys.stderr)
                        except Exception as e:
                            print(f"[WARN] 无法杀掉进程 {pid}: {e}", file=sys.stderr)
    except Exception as e:
        print(f"[WARN] 检查端口占用时出错: {e}", file=sys.stderr)
    
    time.sleep(2)
    print("[INFO] 进程清理完成", file=sys.stderr)

def start_service():
    """启动服务"""
    print("[INFO] 正在启动服务...", file=sys.stderr)
    
    try:
        process = subprocess.Popen([
            "python", "-m", "uvicorn", 
            "services.conversation_service_v2:app", 
            "--host", "0.0.0.0", 
            "--port", "8001"
        ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        
        # 等待服务启动
        for i in range(10):
            time.sleep(1)
            try:
                result = subprocess.run([
                    "curl", "-s", "http://localhost:8001/docs"
                ], capture_output=True, timeout=5)
                if result.returncode == 0:
                    print("[INFO] 服务启动成功", file=sys.stderr)
                    return process
            except Exception:
                pass
            print(f"[INFO] 等待服务启动... ({i+1}/10)", file=sys.stderr)
        
        print("[ERROR] 服务启动失败", file=sys.stderr)
        return None
        
    except Exception as e:
        print(f"[ERROR] 启动服务时出错: {e}", file=sys.stderr)
        return None

async def capture_complete_response():
    """捕获完整的处理流程响应"""
    request_data = {
        "user_id": 1,
        "message": "在平面直角坐标系中，已知点A(0,0)，点B(4,0)，点C(0,3)。\n1. 求三角形ABC的面积\n2. 求三角形ABC的外接圆方程",
        "title": "几何问题测试"
    }
    
    print("[INFO] 开始捕获完整处理流程...", file=sys.stderr)
    print("[INFO] 使用更长的超时时间等待完整处理...", file=sys.stderr)
    
    try:
        # 使用更长的超时时间（10分钟）
        async with httpx.AsyncClient(timeout=600.0) as client:
            async with client.stream("POST", "http://localhost:8001/chat", json=request_data) as response:
                
                # 输出HTTP状态行
                print(f"HTTP/1.1 {response.status_code} {response.reason_phrase}")
                
                # 输出响应头
                for name, value in response.headers.items():
                    print(f"{name}: {value}")
                
                # 空行分隔头部和正文
                print()
                
                # 统计不同类型的消息
                message_count = 0
                step_counts = {}
                
                # 输出响应正文
                try:
                    async for chunk in response.aiter_text():
                        if chunk:
                            print(chunk, end='', flush=True)
                            
                            # 统计消息类型（用于调试）
                            if 'element_extraction:' in chunk:
                                step_counts['element_extraction'] = step_counts.get('element_extraction', 0) + 1
                            elif 'ggb_generation:' in chunk:
                                step_counts['ggb_generation'] = step_counts.get('ggb_generation', 0) + 1
                            elif 'html_generation:' in chunk:
                                step_counts['html_generation'] = step_counts.get('html_generation', 0) + 1
                            elif 'html_optimization:' in chunk:
                                step_counts['html_optimization'] = step_counts.get('html_optimization', 0) + 1
                            elif '"type":"complete"' in chunk:
                                print(f"\n[INFO] 处理完成信号已接收", file=sys.stderr)
                                break
                            
                            message_count += 1
                            
                            # 每1000条消息输出一次进度
                            if message_count % 1000 == 0:
                                print(f"\n[INFO] 已接收 {message_count} 条消息", file=sys.stderr)
                                for step, count in step_counts.items():
                                    print(f"[INFO] {step}: {count} 条", file=sys.stderr)
                
                except Exception as e:
                    if "Broken pipe" not in str(e):
                        print(f"\n[ERROR] 流式处理错误: {e}", file=sys.stderr)
                
                # 输出最终统计
                print(f"\n[INFO] 总共接收 {message_count} 条消息", file=sys.stderr)
                for step, count in step_counts.items():
                    print(f"[INFO] {step}: {count} 条", file=sys.stderr)
                        
    except Exception as e:
        print(f"[ERROR] 捕获响应失败: {e}", file=sys.stderr)
        sys.exit(1)

def main():
    """主函数"""
    service_process = None
    
    try:
        # 1. 清理进程
        cleanup_processes()
        
        # 2. 启动服务
        service_process = start_service()
        if not service_process:
            sys.exit(1)
        
        # 3. 捕获响应
        asyncio.run(capture_complete_response())
        
    except KeyboardInterrupt:
        print("\n[INFO] 用户中断", file=sys.stderr)
    except BrokenPipeError:
        pass
    finally:
        # 清理服务进程
        if service_process:
            print("[INFO] 正在关闭服务...", file=sys.stderr)
            service_process.terminate()
            try:
                service_process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                service_process.kill()
        
        cleanup_processes()

if __name__ == "__main__":
    main() 